
import csv
import os
from pathlib import Path
import pandas as pd


def save_txt(filename, content,mod='w', encoding='utf-8',weather_print = False):
    if weather_print == True:
        print(content)

    with open(filename, mod, encoding=encoding) as file:
        file.write(content)
        file.flush()

def print_save_txt(filename, content,mod='w', encoding='utf-8'):
    print(content)
    save_txt(filename, content, mod = mod, encoding = encoding)

def save_csv(filename, content,mod='w', encoding='utf-8'):
    with open(filename, mod, encoding=encoding,newline='') as file:
        csv_writer = csv.writer(file)
        csv_writer.writerow(content)
        file.flush()

class SaveTxt:
    def __init__(self,fileneme,mod = 'a',encoding='utf-8',weather_print = False):
        self.filename = fileneme
        self.mod = mod
        self.encoding = encoding
        self.weather_print = weather_print

    def save_txt(self,content):
        save_txt(self.filename, content, mod=self.mod, encoding=self.encoding,weather_print= self.weather_print)

class SaveCsv:
    def __init__(self,filename,mod = 'a',encoding='utf-8-sig'):
        self.filename = filename
        self.mod = mod
        self.encoding = encoding
        self.count = 0

    def save_csv(self,content,names = None):
        if names is not None and self.count == 0:
            save_csv(self.filename, names, mod=self.mod, encoding=self.encoding)
            self.count += 1
        else:
            save_csv(self.filename, content, mod=self.mod, encoding=self.encoding)

class SavePicture:
    def __init__(self,folder):
        self.folder = folder
        if Path(self.folder).is_dir():pass #如果文件夹存在就过
        else:os.mkdir(self.folder)  # 创建文件夹

    #这个方法试了下用不了，提示没有.screenshot_as_jpg，还得改成request.get
    def save_pic(self,xpath,filename,driver):
        with open(self.folder + '\\'+ filename, 'wb') as file:
            file.write(driver.find_element_by_xpath(xpath).screenshot_as_jpg)
            file.flush()

class SaveSQL:
    pass

#用于分页保存excel，data_dict{sheet_name:DataFrame}
def save_excel_sheets(data_dict,save_name):
    data_list = data_dict.values()
    sheet_name_list = data_dict.keys()
    writer = pd.ExcelWriter(save_name)
    for data,sheet_name in zip(data_list,sheet_name_list):
        data.to_excel(writer,sheet_name=sheet_name,index=False)
    writer.save()
    writer.close()
